#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time    : 2019-10-08 20:40
@Author  : wangb
@Email   : joker_wb@163.com
@File    : detect_imgs.py
@Software: PyCharm
检测图片中的相关人物信息
"""
from __future__ import division

import datetime
import cv2
from models import *
from utils.yolo_utils import *
from utils.datasets import *
import argparse
import torch
from torch.utils.data import DataLoader
from torch.autograd import Variable

class DetectImgs:
    def __init__(self):
        self.batch_size = 5
        parser = argparse.ArgumentParser()
        parser.add_argument("--image_folder", type=str, default="data/my_samples", help="path to dataset")
        parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
        parser.add_argument("--weights_path", type=str, default="weights/yolov3.weights", help="path to weights file")
        parser.add_argument("--class_path", type=str, default="data/coco.names", help="path to class label file")
        parser.add_argument("--conf_thres", type=float, default=0.8, help="object confidence threshold")
        parser.add_argument("--nms_thres", type=float, default=0.4, help="iou thresshold for non-maximum suppression")
        parser.add_argument("--batch_size", type=int, default=self.batch_size, help="size of the batches")
        parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
        parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
        parser.add_argument("--checkpoint_model", type=str, help="path to checkpoint model")
        self.opt = parser.parse_args()
        print(self.opt)
        self.img_size = self.opt.img_size
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        os.makedirs("output", exist_ok=True)
        # Set up model
        self.model = Darknet(self.opt.model_def, img_size=self.opt.img_size).to(device)
        if self.opt.weights_path.endswith(".weights"):
            # Load darknet weights
            self.model.load_darknet_weights(self.opt.weights_path)
        else:
            # Load checkpoint weights
            self.model.load_state_dict(torch.load(self.opt.weights_path))
        self.model.eval()  # Set in evaluation mode
        self.classes = load_classes(self.opt.class_path)  # Extracts class labels from file
        self.Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
        self.imgs = []  # Stores image paths
        self.img_detections = []  # Stores detections for each image index

    def detect(self, img_lists):
        """
        检测函数
        :param img_lists:
        :return:
        """
        self.imgs = []
        self.img_detections = []
        self.dataloader = DataLoader(
            ImageMemory(img_lists[0:-1], img_size=self.opt.img_size),
            batch_size=self.opt.batch_size,
            shuffle=False,
            num_workers=self.opt.n_cpu,
        )
        for batch_i, (img_indexs, img_tensor) in enumerate(self.dataloader):
            # Configure input
            img_tensor = Variable(img_tensor.type(self.Tensor))

            # Get detections
            with torch.no_grad():
                detections = self.model(img_tensor)
                detections = non_max_suppression(detections, self.opt.conf_thres, self.opt.nms_thres)

            # Save image and detections
            self.imgs.extend(img_indexs)
            self.img_detections.extend(detections)

        # Iterate through images and save plot of detections
        img_show = img_lists[-1]['origin_img']
        rois = img_lists[-1]['rois']
        results = img_lists[-1]
        detect_has_person = []
        for img_i, (img_index, detections) in enumerate(zip(self.imgs, self.img_detections)):
            img = np.array(img_lists[int(img_index)])
            # Draw bounding boxes and labels of detections
            if detections is not None:
                # Rescale boxes to original image
                detections = rescale_boxes(detections, self.opt.img_size, img.shape[:2])
                has_person = False  # 标志位。
                for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
                    if self.classes[int(cls_pred)] != 'person':
                        continue
                    cv2.rectangle(img=img, pt1=(x1, y1), pt2=(x2, y2), color=(0, 0, 255), thickness=2)
                    has_person = True
                    detect_has_person.append(True)

                    break

                if not has_person:
                    detect_has_person.append(False)
            else:
                detect_has_person.append(False)
            """为了显示边框"""
            cv2.rectangle(img=img, pt1=(5, 5), pt2=(img.shape[1] - 10, img.shape[0] - 10), color=(255, 0, 255), thickness=1)
            img_show[rois[img_i][1]:rois[img_i][3] + rois[img_i][1],
            rois[img_i][0]:rois[img_i][2] + rois[img_i][0]] = img
        results['detect_has_person'] = detect_has_person

        """用来检测的时候保存图片的"""

        # hourTime = datetime.datetime.now().strftime('%H:%M:%S')
        # pwd = DetectImgs.mk_dirs('tmp')
        # local_img_path = pwd + '/' + hourTime + '.jpg'
        # cv2.imwrite(local_img_path, img_show)
        results['origin_img'] = img_show
        return results

    @staticmethod
    def mk_dirs(path):
        # 年-月-日 时:分:秒
        nowTime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        # 年-月-日
        dayTime = datetime.datetime.now().strftime('%Y-%m-%d')
        pwd = os.getcwd() + '/' + path + '/' + dayTime
        # 判断文件夹是否已存在
        if not os.path.exists(pwd):
            os.makedirs(pwd)
        return pwd


if __name__ == '__main__':
    di = DetectImgs()
    di.detect()
